উদাহরণসহ ZeRO অপ্টিমাইজার ব্যবহার

ZeRO (Zero Redundancy Optimizer) হল DeepSpeed এর একটি শক্তিশালী ফিচার, যা মেমোরি ব্যবহারে দক্ষতা বাড়াতে এবং বড় মডেল training এর সময় GPU মেমোরির সীমাবদ্ধতা দূর করতে সাহায্য করে। ZeRO এর তিনটি স্তর বা stage আছে:

1. ZeRO Stage 1: Optimizer State Sharding
2. ZeRO Stage 2: Gradient Sharding
3. ZeRO Stage 3: Parameter Sharding

প্রতিটি stage ভিন্ন ভিন্ন optimization প্রদান করে, যা distributed environment এ large-scale মডেল training সহজ করে।

উদাহরণ: ZeRO Optimizer ব্যবহার করে BERT মডেল Training করা

নিচে একটি উদাহরণ দেয়া হলো যেখানে ZeRO Optimizer ব্যবহার করে BERT মডেল train করা হয়েছে।

ধাপ ১: প্রয়োজনীয় প্যাকেজ ইম্পোর্ট করা

DeepSpeed, PyTorch, এবং transformers লাইব্রেরিগুলো ইম্পোর্ট করে শুরু করতে হবে।

import deepspeed
import torch
from transformers import BertForSequenceClassification, BertTokenizer

ধাপ ২: মডেল এবং টোকেনাইজার লোড করা

# BERT মডেল এবং টোকেনাইজার লোড করা
model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

ধাপ ৩: DeepSpeed কনফিগারেশন সেটআপ করা

DeepSpeed এর জন্য একটি কনফিগারেশন সেট করা হবে যেখানে ZeRO Stage 2 এবং Mixed Precision Training ব্যবহার করা হবে।

# deepspeed_config.json ফাইল তৈরি
{
    "train_batch_size": 8,
    "fp16": {
        "enabled": true  # Mixed Precision Training
    },
    "zero_optimization": {
        "stage": 2  # ZeRO Stage 2: Gradient Sharding
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 0.0001
        }
    }
}

কনফিগারেশন ফাইলের প্রতিটি অংশ:

• train_batch_size: Batch size নির্ধারণ করা।
• fp16: FP16 mode enable করে Mixed Precision Training এর মাধ্যমে memory এবং speed efficiency বৃদ্ধি করা।
• zero_optimization: stage 2 নির্বাচন করে gradient sharding এর মাধ্যমে memory efficiency নিশ্চিত করা।
• optimizer: AdamW optimizer ব্যবহার করে learning rate নির্ধারণ করা হয়েছে।

ধাপ ৪: DeepSpeed ইনিশিয়ালাইজ করা এবং মডেল প্রস্তুত করা

DeepSpeed ব্যবহার করে মডেল এবং optimizer ইনিশিয়ালাইজ করা হবে।

# DeepSpeed ইনিশিয়ালাইজ করা
deepspeed_config = "deepspeed_config.json"
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model,
    model_parameters=model.parameters()
)

ধাপ ৫: ডেটাসেট প্রস্তুত করা

ডেটাসেট তৈরি করা এবং সেটিকে DataLoader এ লোড করা। এখানে একটি উদাহরণ হিসেবে ডামি ডেটা ব্যবহার করা হয়েছে।

# ডেটাসেট উদাহরণ (ডামি ডেটা)
texts = ["Hello, how are you?", "DeepSpeed makes training efficient!", "BERT is a powerful model."]
labels = [0, 1, 1]

# Tokenize এবং DataLoader তৈরি করা
inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
inputs["labels"] = torch.tensor(labels)

train_data = torch.utils.data.TensorDataset(inputs["input_ids"], inputs["attention_mask"], inputs["labels"])
train_loader = torch.utils.data.DataLoader(train_data, batch_size=2, shuffle=True)

ধাপ ৬: Training লুপ তৈরি করা

Training লুপ তৈরি করা হবে যেখানে মডেলটি প্রতিটি batch এর জন্য forward এবং backward pass করবে এবং optimizer step নেবে।

# Training লুপ
num_epochs = 3
for epoch in range(num_epochs):
    for batch in train_loader:
        input_ids = batch[0].to(model_engine.local_rank)
        attention_mask = batch[1].to(model_engine.local_rank)
        labels = batch[2].to(model_engine.local_rank)
        
        # Forward pass
        outputs = model_engine(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        
        # Backward pass
        model_engine.backward(loss)
        model_engine.step()  # Optimizer step
        
    print(f"Epoch {epoch + 1} completed. Loss: {loss.item()}")

এই উদাহরণে, DeepSpeed এর ZeRO Optimizer stage 2 ব্যবহার করে gradient partitioning করা হয়েছে, যা মেমোরি ব্যবহারে efficiency বাড়ায়। এছাড়া Mixed Precision Training fp16 enabled করা হয়েছে, যা memory এবং computational efficiency বৃদ্ধি করে।

ZeRO Optimizer এর প্রধান ফিচার এবং উপকারিতা

উপসংহার

ZeRO Optimizer large model training এর ক্ষেত্রে মেমোরি এবং computation efficiency বাড়াতে সহায়ক। এর বিভিন্ন stage, যেমন Stage 1, 2, এবং 3 এর মাধ্যমে optimizer states, gradients, এবং model parameters shard করে memory redundancy কমানো হয়, যা multi-GPU এবং multi-node setup এ বড় মডেল training কে সহজ এবং কার্যকর করে তোলে। DeepSpeed এর সাথে ZeRO Optimizer large-scale distributed training এর জন্য অত্যন্ত গুরুত্বপূর্ণ এবং এটি deep learning এর বড় মডেল training কে আরও scalable এবং efficient করে।